from pymilvus import model

# pip install --upgrade pymilvus
# pip install "pymilvus[model]"
# pip install -q sentence-transformers
# Milvus 通过句子转换器 嵌入函数（SentenceTransformerEmbeddingFunction）类与句子转换器预训练模型集成。
# 该类提供了使用预训练的句子转换器模型对文档和查询进行编码的方法，并将嵌入作为与 Milvus 索引兼容的密集向量返回。

sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(
    model_name='all-MiniLM-L6-v2',  # Specify the model name
    device='cpu'  # Specify the device to use, e.g., 'cpu' or 'cuda:0'
)

docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

docs_embeddings = sentence_transformer_ef.encode_documents(docs)

# Print embeddings
print("Embeddings:", docs_embeddings)
# Print dimension and shape of embeddings
print("Dim:", sentence_transformer_ef.dim, docs_embeddings[0].shape)
